import requests
from .prompt_builder import PromptBuilder
class UserInputHandler:
    def __init__(self, user_input, api_endpoint, model_name, explanations=None, conversation_history=""):
        self.user_input = user_input
        self.api_endpoint = api_endpoint
        self.model_name = model_name
        self.explanations = explanations
        self.conversation_history = conversation_history

    def parse_query(self, df):
        # 使用 PromptBuilder 构建提示词
        prompt_builder = PromptBuilder(df.columns, self.explanations, self.conversation_history, self.user_input)
        prompt = prompt_builder.build_prompt()

        # 构建请求payload
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": False
        }

        try:
            response = requests.post(f"{self.api_endpoint}/generate", json=payload)
            response.raise_for_status()
            
            # 处理响应
            response_data = response.json().get("response").strip()
            
            if response_data.startswith("code"):
                # 提取代码块并返回
                code_block = response_data.split('```python')[1].split('```')[0].strip()
                return code_block
            else:
                # 直接返回自然语言响应
                return response_data

        except requests.exceptions.RequestException as e:
            raise ValueError(f"Error calling model API: {e}")



